import unittest
import pandas as pd
import numpy as np
import os
import glob
import pathlib
import math
import functools   # 排序引用包
from fatigue_algo.fatigue_detector.adaptive_threshold.model.gbdt_for_blink_train import gbdt_infer
from joblib import load

class Headpose_EAR_thres:
    def __init__(self,pitch,yaw,roll,EAR_thres):
        self.pitch = pitch
        self.yaw = yaw
        self.roll = roll
        self.dist = np.round(math.sqrt(math.pow(pitch,2) + math.pow(yaw,2) + math.pow(roll,2)),6)
        self.EAR_thres = EAR_thres

    def __repr__(self):
        return '(%f,%f,%f,%f,%f)' % (self.dist, self.pitch, self.yaw, self.roll, self.EAR_thres)

#对Headpose_EAR_thres对象按属性进行排序
def my_cmp(self,object):
    """
    自定义排序方法
    :param object:  Headpose_EAR_thres对象
    :return: 先按照dist从小到大，若等则按pitch从小到大，若等则按照yaw从小到大，若等则按照roll从小到大
    """
    if self.dist < object.dist:
        return -1
    elif self.dist > object.dist:
        return 1
    else:
        if self.pitch < object.pitch:
            return -1
        elif self.pitch > object.pitch:
            return 1
        else:
            if self.yaw < object.yaw:
                return -1
            elif self.roll > object.roll:
                return 1
            else:
                if self.roll < object.roll:
                    return -1
                elif self.roll > object.roll:
                    return 1
                else:
                    return 0

#读取数据集中指定的标注文件,并将所有标注文件转化成一个新的数据集
def convert_dataset_annotation(in_Dir, out_path, func):
    '''
    :param in_Dir: 自适应EAR数据集文件夹路径，文件结构如下：
        |_up
            |_center
            |_lower_left
            |_lower_right
            |_upper_left
            |_upper_right
        |_right
        |_left
        |_front
        |_down
    :param out_path:
    :param func: 要按文件夹处理的函数，包括get_EAR_thres_means和
    :return:
    '''
    # 获取一级文件夹
    for dir in pathlib.Path(in_Dir).iterdir():
        # 获取二级文件夹
        if(os.path.isdir(str(dir)) == False): #非文件夹，则返回
            continue
        for sub_dir in pathlib.Path(str(dir)).iterdir():
            # 获取标注文件
            totalPaths = glob.glob(os.path.join(str(sub_dir),"*.csv"))
            for path in totalPaths:
                print("reading " + path)
                temp = func(path,out_path)
                print(temp)

#将每个视频的标注文件转化成一个样本，每个属性分别对应pitch,yaw,roll,EAR_threshold的均值
def get_EAR_thres_means(in_path, out_path):
    '''
    :param in_path: 源数据集标注文件路径
    :param out_path: 要写入的文件路径
    :return:
    '''
    df = pd.read_csv(in_path)  #源数据集标注文件
    #先判断out_path文件是否存在，如果不存在则创建空表
    if(os.path.exists(out_path)):
        res_df = pd.read_csv(out_path)
    else:
        res_df = pd.DataFrame(columns=['pitch','yaw','roll','EAR_thres'])

    temp = []
    #pitch,yaw,roll的均值
    for column in df.columns[1:-2]:
        mean = np.mean(df[column].values)
        temp.append(np.round(mean,6))

    #EAR_threshold均值
    EARs = df[df.columns[-2]].values   #EAR值
    blink_state = df[df.columns[-1]].values  #是否眨眼
    EAR_thresholds = []
    # for i in range(1,len(EARs) - 1):
    #     # 闭眼（01边界）
    #     if(blink_state[i - 1] == 0 and blink_state[i] == 1):
    #         EAR_thresholds.append(EARs[i])
    #     # 睁眼（10边界）
    #     elif(blink_state[i] == 1 and blink_state[i + 1] == 0):
    #         EAR_thresholds.append(EARs[i])
    for i in range(0,len(EARs)):
        # 眨眼
        if(blink_state[i] == 1):
            EAR_thresholds.append(EARs[i])
    # res_EAR_thresholds = np.round(np.mean(EAR_thresholds),6)
    # res_EAR_thresholds = np.round(np.max(EAR_thresholds),6)
    res_EAR_thresholds = np.percentile(EAR_thresholds, 99)
    temp.append(res_EAR_thresholds)

    res_df.loc[len(res_df)] = temp
    res_df.to_csv(out_path,index=False)
    return temp

#将文件夹下所有标注文件中的眨眼样本EAR值修改EAR的上85分位数，整合成一个训练集
def get_blink_train_set(in_path, out_path):
    '''
    :param in_path: 源数据集标注文件路径
    :param out_path: 要写入的文件路径
    :return:
    '''
    df = pd.read_csv(in_path)  # 源数据集标注文件
    # 先判断out_path文件是否存在，如果不存在则创建空表
    if (os.path.exists(out_path)):
        res_df = pd.read_csv(out_path)
    else:
        res_df = pd.DataFrame(columns=['pitch', 'yaw', 'roll', 'EAR_thres'])

    # EAR_threshold均值
    EARs = df[df.columns[-2]].values  # EAR值
    blink_state = df[df.columns[-1]].values  # 是否眨眼
    EAR_thresholds = []
    #获取EARs的85分位数
    for i in range(0, len(EARs)):
        # 眨眼
        if (blink_state[i] == 1):
            EAR_thresholds.append(EARs[i])
    res_EAR_thresholds = np.percentile(EAR_thresholds, 85)

    #眨眼样本生成
    columns = df.columns
    for i in range(0, len(EARs)):
        # 眨眼
        if (blink_state[i] == 1):
            temp = []
            temp.extend([np.round(df.loc[i][columns[1]],6),np.round(df.loc[i][columns[2]],6),np.round(df.loc[i][columns[3]],6)])
            temp.append(res_EAR_thresholds)
            res_df.loc[len(res_df)] = temp

    res_df.to_csv(out_path, index=False)

#获取正脸最大的EAR，可用在其他人的人脸校准上
def get_frontFace_max_EAR(in_path):
    df = pd.read_csv(in_path)  # 源数据集标注文件
    EARs = df['EAR'].values #获取EARs
    return np.max(EARs)

#根据给定的pitch,yaw,roll找出最相似的姿态，得到此时的EAR_thres
def get_EAR_thres_withHeadPose(in_path, pitch, yaw, roll):
    '''
    :param in_path: 自适应阈值文件
    :param pitch: 当前头部姿态的俯仰角
    :param yaw: 当前头部姿态的偏航角
    :param roll: 当前头部姿态的旋转角
    :return:
    '''
    df = pd.read_csv(in_path)
    columns = df.columns
    objects = []
    for i in range(0,len(df)):
        delta_pitch = np.abs(pitch - df.loc[i][columns[0]])
        delta_yaw = np.abs(yaw - df.loc[i][columns[1]])
        delta_roll = np.abs(roll - df.loc[i][columns[2]])
        EAR = df.loc[i][columns[3]]
        object = Headpose_EAR_thres(delta_pitch,delta_yaw,delta_roll,EAR)
        objects.append(object)

    #对objects按自定义排序方式进行排序
    objects = sorted(objects,key=functools.cmp_to_key(my_cmp))  #https://blog.csdn.net/chen1231985111/article/details/121426195
    EAR_thres = objects[0].EAR_thres  #获取最相似的头部姿态的EAR_thres
    return EAR_thres

#通过相似的姿态获得的EAR_thres只适合A受试者，对于B受试者需要通过传入正脸的EAR进行校正
def calibrate_EAR_thres(EAR_A,EAR_B,EAR_thres):
    return EAR_B * EAR_thres / EAR_A

class test(unittest.TestCase):

    #测试get_EAR_thres_means
    def test_get_EAR_thres_means(self):
        in_path = "F:/fatigue_recognition/dataset/adaptive_thres/up/center/2022_08_19_08_08_29.csv"
        out_path = "./data/adaptive_threshold.csv"
        temp = get_EAR_thres_means(in_path,out_path)
        print(temp)

    #测试基于convert_dataset_annotation的get_EAR_thres_means方法
    def test_convert_dataset_annotation_get_EAR(self):
        in_Dir = "F:/fatigue_recognition/dataset/adaptive_thres/"
        out_path = "./data/adaptive_threshold_99.csv"
        convert_dataset_annotation(in_Dir,out_path,get_EAR_thres_means)

    # 测试get_blink_train_set方法
    def test_get_blink_train_set(self):
        in_path = "F:/fatigue_recognition/dataset/adaptive_thres/up/center/2022_08_19_08_08_29.csv"
        out_path = "./data/blink_train.csv"
        temp = get_blink_train_set(in_path,out_path)
        print(temp)

    # 测试基于convert_dataset_annotation的get_blink_train_set方法
    def test_convert_dataset_annotation_get_blink_train(self):
        in_Dir = "F:/fatigue_recognition/dataset/adaptive_thres/"
        out_path = "./data/blink_train_85.csv"
        convert_dataset_annotation(in_Dir,out_path,get_blink_train_set)

    # 测试test_get_frontFace_max_EAR
    def test_get_frontFace_max_EAR(self):
        in_path = "F:/fatigue_recognition/dataset/adaptive_thres/up/center/2022_08_19_08_08_29.csv"
        EAR_max = get_frontFace_max_EAR(in_path)
        print(EAR_max)

    # 测试get_EAR_thres_withHeadPose
    def test_get_EAR_thres_withHeadPose(self):
        in_path = "./data/adaptive_threshold.csv"
        pitch, yaw, roll = 0,0,50
        EAR_thres = get_EAR_thres_withHeadPose(in_path,pitch, yaw, roll)
        print(EAR_thres)

    # 获取相似度最大的头部姿态对应的EAR值，测试calibrate_EAR_thres
    def test_calibrate_EAR_thres(self):

        in_path = "F:/fatigue_recognition/dataset/adaptive_thres/up/center/2022_08_19_08_08_29.csv"
        EAR_max = get_frontFace_max_EAR(in_path)

        in_path = "./data/adaptive_threshold.csv"
        pitch, yaw, roll = 0, 0, 0
        EAR_thres = get_EAR_thres_withHeadPose(in_path, pitch, yaw, roll)

        EAR_A, EAR_B = EAR_max, np.float64(0.25)
        EAR_thres = calibrate_EAR_thres(EAR_A,EAR_B,EAR_thres)
        print(EAR_thres)

    # 利用gbdt预测此时头部姿态的EAR值，测试calibrate_EAR_thres
    def test_calibrate_EAR_thres_with_GBDT(self):

        in_path = "F:/fatigue_recognition/dataset/adaptive_thres/up/center/2022_08_19_08_08_29.csv"
        EAR_max = get_frontFace_max_EAR(in_path)

        pitch, yaw, roll = 0, 0, 0
        gbr = load("model/blink_EAR_thres_gbdt.joblib")
        EAR_thres = gbdt_infer(gbr,pitch, yaw, roll)
        print(f"stage1 : {EAR_thres}")

        EAR_A, EAR_B = EAR_max, np.float64(0.25)
        EAR_thres = calibrate_EAR_thres(EAR_A,EAR_B,EAR_thres)
        print(f"stage2 : {EAR_thres}")

